/*
 *  Copyright 2020-2025 the original author or authors.
 *  You cannot use this file unless authorized by the author.
 */

package org.ipig.computing.spark.connector;

import com.mongodb.client.model.IndexModel;
import com.mongodb.client.model.IndexOptions;
import com.mongodb.client.model.WriteModel;
import com.mongodb.spark.MongoConnector;
import com.mongodb.spark.config.WriteConfig;
import lombok.extern.slf4j.Slf4j;
import org.bson.Document;

import java.util.ArrayList;
import java.util.List;

/**
 * SparkMongoOperators
 *
 * @author <a href="mailto:comchnts@163.com">chinats</a>
 * @since 1.0
 */
@Slf4j
public class SparkMongoOperators {
    /**
     * dropCollection
     *
     * @param create 集合信息
     */
    public static boolean dropCollection(WriteConfig create) {
        return MongoConnector.apply(create.asOptions()).withCollectionDo(create, Document.class, (v1) -> {
            log.info("删除表: " + v1.getNamespace());
            v1.drop();
            return true;
        });
    }

    /**
     * saveBatchRes
     *
     * @param writeConfig    数据db信息
     * @param mongoConnector mongodb连接器
     * @param bulkModels     批量model
     */
    public static boolean bulkWrite(WriteConfig writeConfig, MongoConnector mongoConnector,
                                    List<WriteModel<Document>> bulkModels) {
        return mongoConnector.withCollectionDo(writeConfig, Document.class, (v1) -> {
            v1.bulkWrite(bulkModels);
            return true;
        });
    }


    /**
     * createIndexes
     *
     * @param writeConf
     * @param indexList
     */
    public static boolean createIndexes(WriteConfig writeConf, List<Document> indexList) {
        boolean flag = false;
        if (indexList != null && !indexList.isEmpty()) {
            flag = MongoConnector.apply(writeConf.asOptions()).withCollectionDo(writeConf, Document.class, (v1) -> {
                List<Document> listIndexes = new ArrayList<>();
                v1.listIndexes().into(listIndexes);
                List<IndexModel> indexs = new ArrayList<>();
                int i = 0;
                if (v1.countDocuments() == 0 && listIndexes.size() < 1) {
                    for (Document doc : indexList) {
                        indexs.add(new IndexModel(doc, new IndexOptions().name("index_" + (i++))));
                    }
                    // 创建索引
                    try {
                        v1.createIndexes(indexs);
                    } catch (Exception e) {
                        log.error(e.getMessage());
                    }
                }
                return true;
            });
        }
        return flag;
    }
}
